import os
import sys
sys.path.append('..')
import cv2
import pandas as pd
from pathlib import Path
import glob
from utils.json import json_to_image

def get_mask(img_path, dst_path):
    code = Path(img_path).parent.name

    img = cv2.imread(img_path, 0)

    json_path = Path(img_path).with_suffix('.json')
    if not os.path.exists(json_path):
        return

    try:
        mask = json_to_image(img_shape=img.shape, json_path=json_path)
        mask *= 255

        mask_name = Path(img_path).with_suffix('.png').name
        mask_dst_path = os.path.join(dst_path, code, mask_name)
        os.makedirs(os.path.join(dst_path, code), exist_ok=True)

        cv2.imwrite(mask_dst_path, mask)
    except:
        return 

def main():
    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False

    path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6007/all_0409'
    dst_path = r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/t7/T6007/all_0409_mask'
    df = pd.DataFrame()
    df['img_path'] = glob.glob(os.path.join(path, "*/*.jpg"))
    print(len(df))

    if is_pandarallel:
        df.parallel_apply(lambda x:get_mask(img_path=x['img_path'], dst_path=dst_path), axis=1)
    else:
        df.apply(lambda x:get_mask(img_path=x['img_path'], dst_path=dst_path), axis=1)

    print()


if __name__=='__main__':
    main()
